from manimlib import *

__curr_dir = os.path.dirname(__file__)
if __curr_dir not in sys.path:
    sys.path.append(__curr_dir)


turns = 20

def Spring3D(start, end, axis):
    """
    创建弹簧，包括首尾节点（球体）和长度标签
    """

    start = np.array(start)
    end = np.array(end)

    global_length = np.linalg.norm(end - start)

    start = axis.c2p(start[0], start[1], start[2])
    end = axis.c2p(end[0], end[1], end[2])

    # 横截面半径
    radius = global_length * 0.05

    curve = mobject.functions.ParametricCurve(
        lambda t: axis.c2p(
            t, radius * np.cos(2 * np.pi * turns * t), radius * np.sin(2 * np.pi * turns * t)  # x坐标  # y坐标  # z坐标
        ),
        t_range=[0, 1, 0.005],  # 绘制步长0.001
        color=YELLOW,
    )
    curve.set_style(stroke_width=10)

    # 将单元弹簧映射到真实坐标上
    curve.shift(start)  # 将弹簧移动到起始位置
    curve.stretch(global_length, 0, about_point=start)  # 根据向量的长度来缩放弹簧

    # # 在弹簧的两端添加两个球体
    # ball_start = Dot(radius=radius * 0.2, color=RED).move_to(start)

    return curve


class SpringTruss(Scene):
    def construct(self):
        axes = ThreeDAxes()
        frame = self.camera.frame
        frame.set_height(12)

        # ========================================================
        # ========================================================
        # ========================================================
        # 创建弹簧
        sprint_start = [0, 0, 0]
        sprint_end = [10, 0, 0]
        spring = Spring3D(sprint_start, sprint_end, axes)
        self.add(spring)
        
        truss_node_scale = 2

        nodes = Group(
            # Sphere(color=RED).move_to(np.array([0, 0, 0])),
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([0 * truss_node_scale, 0 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 0
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([0 * truss_node_scale, 1 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 1
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([0 * truss_node_scale, 2 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 2
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([0 * truss_node_scale, 3 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 3
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([0 * truss_node_scale, 4 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 4
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([1 * truss_node_scale, 4 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 5
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([2 * truss_node_scale, 4 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 6
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([3 * truss_node_scale, 4 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 7
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([3 * truss_node_scale, 3 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 8
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([2 * truss_node_scale, 3 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 9
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([1 * truss_node_scale, 3 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 0
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([1 * truss_node_scale, 2 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 11
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([1 * truss_node_scale, 1 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 12
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([1 * truss_node_scale, 0 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 13
            Sphere(color=RED, radius=0.1)
            .move_to(np.array([4 * truss_node_scale, 3 * truss_node_scale, 0]))
            .shift([-4, -5, 0]),  # 14
        )

        elements = Group(
            Line(start=nodes[0], end=nodes[1], color=YELLOW),
            Line(start=nodes[1], end=nodes[2], color=YELLOW),
            Line(start=nodes[2], end=nodes[3], color=YELLOW),
            Line(start=nodes[3], end=nodes[4], color=YELLOW),
            Line(start=nodes[4], end=nodes[5], color=YELLOW),
            Line(start=nodes[5], end=nodes[6], color=YELLOW),
            Line(start=nodes[6], end=nodes[7], color=YELLOW),
            Line(start=nodes[7], end=nodes[8], color=YELLOW),
            Line(start=nodes[8], end=nodes[9], color=YELLOW),
            Line(start=nodes[9], end=nodes[10], color=YELLOW),
            Line(start=nodes[10], end=nodes[11], color=YELLOW),
            Line(start=nodes[11], end=nodes[12], color=YELLOW),
            Line(start=nodes[12], end=nodes[13], color=YELLOW),
            Line(start=nodes[13], end=nodes[0], color=YELLOW),
            Line(start=nodes[13], end=nodes[1], color=YELLOW),
            Line(start=nodes[12], end=nodes[1], color=YELLOW),
            Line(start=nodes[12], end=nodes[2], color=YELLOW),
            Line(start=nodes[11], end=nodes[2], color=YELLOW),
            Line(start=nodes[11], end=nodes[3], color=YELLOW),
            Line(start=nodes[10], end=nodes[3], color=YELLOW),
            Line(start=nodes[10], end=nodes[4], color=YELLOW),
            Line(start=nodes[10], end=nodes[5], color=YELLOW),
            Line(start=nodes[9], end=nodes[5], color=YELLOW),
            Line(start=nodes[9], end=nodes[6], color=YELLOW),
            Line(start=nodes[8], end=nodes[6], color=YELLOW),
            Line(start=nodes[7], end=nodes[14], color=YELLOW),
            Line(start=nodes[8], end=nodes[14], color=YELLOW),
        )

        self.play(
            spring.animate.scale(0.15).move_to(elements[4].get_center()),
        )
        self.remove(spring)
        self.add(elements[4])
        Truss_force_arrow = Arrow(
            start=nodes[14],
            end=nodes[14].get_center() - np.array([0, 2, 0]),
        )
        Truss_force_arrow.set_style(stroke_color=RED)
        Truss_force_text = Text("F", color=RED).next_to(Truss_force_arrow, DOWN)
        Truss_boundary = Line(start=nodes[0].shift(LEFT), end=nodes[13].shift(RIGHT), color=GREY)
        Truss_boundary_text = Text("Boundary", color=GREY).next_to(Truss_boundary, DOWN)

        Truss_model = Group(
            nodes,
            elements,
            Truss_force_arrow,
            Truss_force_text,
            Truss_boundary,
            Truss_boundary_text,
        )
        self.add(Truss_model)
        self.wait(2)
        self.play(Truss_model.animate.scale(0.8).shift(LEFT * 5))
        self.wait(2)

        Truss_global_equal = Tex(
            r"\begin{bmatrix} F_{1} \\ F_{2} \\ \vdots \\ F_{n} \end{bmatrix} =",
            r"\begin{bmatrix} \frac{EA}{L_{11}} & \frac{EA}{L_{12}} & \cdots & \frac{EA}{L_{1n}} \\ \frac{EA}{L_{21}} & \frac{EA}{L_{22}} & \cdots & \frac{EA}{L_{2n}} \\ \vdots & \vdots & \ddots & \vdots \\ \frac{EA}{L_{n1}} & \frac{EA}{L_{n2}} & \cdots & \frac{EA}{L_{nn}} \end{bmatrix}",
            r"\begin{bmatrix} u_{1} \\ u_{2} \\ \vdots \\ u_{n} \end{bmatrix}",
        )
        Truss_global_equal.shift(RIGHT * 4).scale(1.2)
        self.play(ShowCreation(Truss_global_equal))
        self.wait(2)